
from tensorflow.keras import backend
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Activation
from tensorflow.keras.models import Sequential


class Shallow_Net:
    @staticmethod
    def build(width, height, depth, classes) -> Sequential:
        """
        构建网络函数

        :param width: 矩阵列数
        :param height: 矩阵行数
        :param depth: 图像通道数
        :param classes: 图像分类数
        :return: 模型对象
        """
        # 此处为构建输入数据的深度优先集
        model = Sequential()
        input_shape = (height, width, depth)

        if backend.image_data_format() == "channels_first":
            input_shape = (depth, height, width)

        # 定义第一层卷积（conv=>Relu）
        model.add(Conv2D(32, (3, 3), padding="same",
                         input_shape=input_shape))
        model.add(Activation("relu"))

        # 将输入数据平展
        model.add(Flatten())
        model.add(Dense(classes))
        model.add(Activation("softmax"))

        return model
